"""This program convert gromacs itp file into dff ppf file.
Only OPLS was supported now!
By madlee@2009-1-5"""

import sys

USAGE = """usage: gmx2ppf nonbond.itp valence.itp output.ppf"""

KJ_TO_KCAL = 1.0 / 4.184

def parse_nonbond(file):
    SIGMA_TO_R0 = 10 * pow(2.0, 1.0/6.0)
    for line in file:
        if line.startswith('[ atomtypes ]'):
            break
    nonbond = []
    eqt = []
    for line in file:
        line = line.strip()
        if line.startswith(';'):
            pass
        else:
            tokens = line.split()
            if len(tokens) >= 8:
                eqt.append((tokens[0], tokens[1]))
                atc = "ATC: " + tokens[0] + ": " + tokens[4]
                nonbond.append(atc)
                atype = "ATYPE: " + tokens[0] + ": " + tokens[2] + ", " + tokens[3]
                nonbond.append(atype)
                r0 = float(tokens[6]) * SIGMA_TO_R0
                e0 = float(tokens[7]) * KJ_TO_KCAL
                n126 = "N12_6: %s: %12.4f, %12.4f" % (tokens[0], r0, e0)
                nonbond.append(n126)
    return (nonbond, eqt)

def saveeqt(eqt, filename):
    pos = filename.rfind('.')
    if pos != -1:
        filename = filename[:pos] + ".eqt"
    else:
        filename += ".eqt"

    with open(filename, "wb") as file:
        file.write("#AAT                :	NB	BINC	Bond	A/C	A/S	T/C	T/S	O/C	O/S\n")
        for i in eqt:
            file.write(i[0] + ":\t" + i[0] + "\t")
            file.write("\t".join([i[1]] * 8))
            file.write("\n")

def parse_bonds(file):
    result = []
    next_token = ""
    for i in file:
        if i.startswith(';') or i.startswith('#'):
            pass
        elif i.startswith("["):
            next_token = i
            break
        else:
            tokens = i.strip().split()
            if len(tokens) >= 5:
                r0 = float(tokens[3]) * 10
                k0 = float(tokens[4]) / 200 * KJ_TO_KCAL
                bharm = "BHARM: %s, %s: %12.4f, %12.4f" % (tokens[0], tokens[1], r0, k0)
                result.append(bharm)

    return (result, next_token)

def parse_angles(file):
    result = []
    next_token = ""
    for i in file:
        if i.startswith(';') or i.startswith('#'):
            pass
        elif i.startswith("["):
            next_token = i
            break
        else:
            tokens = i.strip().split()
            if len(tokens) >= 5:
                t0 = float(tokens[4])
                k0 = float(tokens[5]) * KJ_TO_KCAL
                aharm = "AHARM: %s, %s, %s: %12.4f, %12.4f" % (tokens[0], tokens[1], tokens[2], t0, k0)
                result.append(aharm)

    return (result, next_token)

def parse_torsions(file):
    result = []
    next_token = ""
    for i in file:
        if i.startswith(';') or i.startswith('#'):
            pass
        elif i.startswith("["):
            next_token = i
            break
        else:
            tokens = i.strip().split()
            if len(tokens) >= 5:
                pass

    return (result, next_token)

def parse_valence(file):
    result = []
    nexttoken = "1"
    while nexttoken:
        if nexttoken.startswith("[ bondtypes ]"):
            (items, nexttoken) = parse_bonds(file)
            result += items
        elif nexttoken.startswith("[ angletypes ]"):
            (items, nexttoken) = parse_angles(file)
            result += items
        elif nexttoken.startswith("[ dihedraltypes ]"):
            (items, nexttoken) = parse_torsions(file)
            result += items
        else:
            nexttoken = file.next()
        
    return result
            
    
    
def saveppf(ppf, filename):
    with open(filename, "wb") as file:
        file.write('#DFF:PPF\n#PROTOCOL = OPLS\n')
        for i in ppf:
            file.write(i + '\n')

if __name__ == '__main__':
    if len(sys.argv) != 4:
        print USAGE
    else:
        with open(sys.argv[1]) as file:
            ppf, eqt = parse_nonbond(file)

        saveeqt(eqt, sys.argv[3])

        with open(sys.argv[2]) as file:
            ppf += parse_valence(file)
            
        ppf.sort()
        saveppf(ppf, sys.argv[3])
